#!/usr/bin/env python3

import rospy
import os
from sensor_msgs.msg import NavSatFix
from nav_msgs.msg import Odometry
from std_msgs.msg import String
import datetime

class DataLoggerNode:
    def __init__(self):
        rospy.init_node('data_logger_node')

        # 获取脚本所在目录
        script_dir = os.path.dirname(os.path.realpath(__file__))

        # 文件名 (使用日期时间)
        timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
        self.navsat_filename = os.path.join(script_dir, f"navsat_{timestamp}.txt")
        self.odom_filename = os.path.join(script_dir, f"odom_{timestamp}.txt")
        self.raw_filename = os.path.join(script_dir, f"raw_{timestamp}.txt")

        # 打开文件
        self.navsat_file = open(self.navsat_filename, "w")
        self.odom_file = open(self.odom_filename, "w")
        self.raw_file = open(self.raw_filename, "w")

        # 写入文件头
        self.navsat_file.write("timestamp,latitude,longitude,altitude,status\n")
        self.odom_file.write("timestamp,x,y,yaw,vx,vy\n")
        self.raw_file.write("timestamp,raw_nmea_sentence\n")


        # 订阅者
        self.navsat_sub = rospy.Subscriber('/rtk_reader_node/navsat', NavSatFix, self.navsat_callback)
        self.odom_sub = rospy.Subscriber('/rtk_odometry_node/odom', Odometry, self.odom_callback)
        self.raw_sub = rospy.Subscriber('/rtk_reader_node/raw', String, self.raw_callback)

        rospy.loginfo(f"NavSatFix data will be saved to: {self.navsat_filename}")
        rospy.loginfo(f"Odometry data will be saved to: {self.odom_filename}")
        rospy.loginfo(f"Raw NMEA data will be saved to: {self.raw_filename}")


    def navsat_callback(self, msg):
        timestamp = msg.header.stamp.to_sec()
        line = f"{timestamp},{msg.latitude},{msg.longitude},{msg.altitude},{msg.status.status}\n"
        self.navsat_file.write(line)
        self.navsat_file.flush()

    def odom_callback(self, msg):
        timestamp = msg.header.stamp.to_sec()
        x = msg.pose.pose.position.x
        y = msg.pose.pose.position.y
        quaternion = (
            msg.pose.pose.orientation.x,
            msg.pose.pose.orientation.y,
            msg.pose.pose.orientation.z,
            msg.pose.pose.orientation.w
        )
        from tf.transformations import euler_from_quaternion
        _, _, yaw = euler_from_quaternion(quaternion)

        vx = msg.twist.twist.linear.x
        vy = msg.twist.twist.linear.y
        line = f"{timestamp},{x},{y},{yaw},{vx},{vy}\n"
        self.odom_file.write(line)
        self.odom_file.flush()

    def raw_callback(self, msg):
        timestamp = rospy.Time.now().to_sec()
        line = f"{timestamp},{msg.data}\n"
        self.raw_file.write(line)
        self.raw_file.flush()

    def on_shutdown(self):
        self.navsat_file.close()
        self.odom_file.close()
        self.raw_file.close()
        rospy.loginfo("Files closed.")

if __name__ == '__main__':
    try:
        logger = DataLoggerNode()
        rospy.on_shutdown(logger.on_shutdown)
        rospy.spin()
    except rospy.ROSInterruptException:
        pass